import sys
sys.path.append('../face_recog/insightface/deploy')
sys.path.append('../face_recog/insightface/src/common')
from imutils import paths
from face_recog.insightface.deploy import face_model
import numpy as np
from pathlib import Path
import argparse
import pickle
import cv2
import os
from face_embedding.utils import argsutils
import faiss
import time
import threading
import json
import re
import math

from utils import dictProperties


args = argsutils.Facemodel(dictProperties).args

class faiss_index():
    def __init__(self,args, datas):
        self.args = args
        self.d = 512
        self.start_time = time.time()
        self.user_codes = []
        self.embeddings = []
        for user_code, embedd in datas:
            self.user_codes.append(user_code)
            self.embeddings.append(embedd)
        self.parameters_init()

    def renew(self, datas):
        self.user_codes = []
        self.embeddings = []
        for user_code, embedd in datas:
            self.user_codes.append(user_code)
            print('---------renew', len(embedd))
            self.embeddings.append(embedd)

        self.parameters_init()

    def parameters_init(self):
        self.ids = []
        self.index = faiss.IndexFlatL2(self.d)
        # print(self.embeddings)
        data_list = []
        index = 0
        for embed in self.embeddings:
            # print('-----parameters_init')
            # print(type(embed), len(embed))
            for e in embed:
                self.ids.append(self.user_codes[index])
                data_list.append(e)
                # print('---parameters_init:', len(e))
            index += 1
        d = np.array(data_list, dtype=np.float32)
        # print(d.shape)
        # print(d.dtype)
        print('ids len:', len(self.ids), len(data_list))
        self.index.add(d)

    def add_embedding(self, user_code, embedding):
        data_list = []
        # index_tem = faiss.IndexFlatL2(self.d)
        # print('add_embedding')
        # print('user_code:', user_code)
        # print('embedding len:', len(embedding))
        for embed in embedding:
            self.ids.append(user_code)
            data_list.append(embed)

        self.index.add(np.array(embedding, dtype=np.float32))

    def get_index(self):
        return self.index

    def get_ids(self):
        return self.ids

class faiss_group():
    def __init__(self, args):
        self.faiss_dict = {}
        self.args = args

    def refresh_faiss(self, group_id, datas):
        faiss_obj = self.faiss_dict.get(group_id)
        if faiss_obj is None:
            faiss_ind = faiss_index(self.args, datas)
            self.faiss_dict[group_id] = faiss_ind
        else:
            faiss_obj.renew(datas)

    def get_faiss(self, group_id):
        faiss_obj = self.faiss_dict.get(group_id)
        return faiss_obj

    def add_faiss_embedding(self, group_id, face_id, embedding):
        faiss_obj = self.faiss_dict.get(group_id)
        if faiss_obj == None:
            datas = [(face_id, embedding)]
            faiss_ind = faiss_index(self.args, datas)
            self.faiss_dict[group_id] = faiss_ind
        else:
            faiss_obj.add_embedding(face_id, embedding)

    def delete_faiss_embedding(self, group_id):
        self.faiss_dict.pop(group_id, None)
        print('delete db:', group_id)


    def init_faiss(self, group_data):
        for group_id in group_data.keys():
            datas = group_data.get(group_id)
            faiss_ind = faiss_index(self.args, datas)
            self.faiss_dict[group_id] = faiss_ind

faiss_dict = faiss_group(args)

def init_faiss(db):
    sql_s = "select id, group_id, face_id, feature_data, ext_id, create_time from t_home_user_face"
    data_dict = {}
    result = db.execute_query(sql_s, ())
    if result:
        for user_info in result:
            face_id = user_info[2]
            group_id = user_info[1]
            embedding = user_info[3]

            res = []
            # for emb_lst in json.loads(embedding):
            #     # lst = re.findall('-?\d+.\d+', emb_lst)
            #     # tem = [float(i) for i in lst]
            #     # res.append(np.array(tem))

            res.append(np.array(json.loads(embedding), dtype=np.float32))

            data_list = data_dict.get(group_id)
            if data_list is None:
                data_list = []
                data_dict[group_id] = data_list
            data_list.append((face_id, res))
        faiss_dict.init_faiss(data_dict)

# 新增faiss_dict
def add_faiss_embedding(group_id, face_id, embedding):
    faiss_dict.add_faiss_embedding(group_id, face_id, embedding)

def delete_faiss_embedding(group_id):
    faiss_dict.delete_faiss_embedding(group_id)

# 更新faiss
def refresh_faiss(group_id, db):
    def thread_run(faiss_dict, group_id, db):
        sql_s = "select id, group_id, face_id, feature_data, ext_id, create_time from t_home_user_face where group_id='{}'".format(group_id)
        result = db.execute_query(sql_s, ())
        data_list = []
        if result:
            for user_info in result:
                face_id = user_info[2]
                group_id = user_info[1]
                embedding = user_info[3]

                res = []
                # for emb_lst in json.loads(embedding):
                #     # lst = re.findall('-?\d+.\d+', emb_lst)
                #     # tem = [float(i) for i in lst]
                #     # res.append(np.array(tem))

                res.append(np.array(json.loads(embedding), dtype=np.float32))

                data_list.append((face_id, res))
            faiss_dict.refresh_faiss(group_id, data_list)

    t = threading.Thread(target=thread_run, args=(faiss_dict, group_id, db))
    t.start()

def faiss_search(groupId, img, embedding_model):
    # frame = cv2.imread(img_dir)
    # print(faiss_ind.start_time)
    faiss_obj = faiss_dict.get_faiss(groupId)
    if faiss_obj is None:
        return (90009, '该用户组未注册'), None, None, None
    index = faiss_obj.get_index()
    ids = faiss_obj.get_ids()

    # image = cv2.resize(frame, (112, 112))
    frame = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    frame = np.transpose(frame, (2, 0, 1))
    # Get the face embedding vector
    face_embedding = embedding_model.get_feature(frame)

    res,I,D=idx_img(index, ids, face_embedding)
    print(' user_id = %s'%(res))
    return None, res,I,D


def faiss_search_embedding(groupId, face_embedding, top_n, threshold):
    # frame = cv2.imread(img_dir)
    # print(faiss_ind.start_time)
    faiss_obj = faiss_dict.get_faiss(groupId)
    if faiss_obj is None:
        return (90009, '该用户未注册'), None, None, None
    index = faiss_obj.get_index()
    ids = faiss_obj.get_ids()

    res, I, D = idx_img(index, ids, face_embedding, top_n, threshold)
    print(' user_id = %s' % (res))
    return None, res, I, D

def idx_img(faiss_idx, id_lst, face_embedding, top_n, threshold):
    # print(id_lst)
    embedding = np.array(face_embedding, dtype=np.float32)
    print(embedding.shape)
    D, I = faiss_idx.search(np.expand_dims(embedding, 0), top_n)
    I_val = I.tolist()
    D_val = D.tolist()
    # print(I)
    # 这边设定阈值
    print(len(I), len(D), len(id_lst))
    print(I, D)
    print(id_lst)
    #if D[0][0] > threshold:
    #    return None,I_val,D_val
    # print(name_lst)
    #if I[0][0] > len(id_lst):
    #    return None,I_val,D_val
    #ids = []
    #for i in range(top_n):
    #    id_val = I[0][i]
    #    if id_val >= 0:
    #        ids.append(id_lst[id_val])

    dis_thred=2-2*threshold
    if D[0][0] > dis_thred:
        return None,I_val,D_val
    print(I,D)
    # print(name_lst)
    if I[0][0] > len(id_lst):
        return None,I_val,D_val
    ids = []
    for i in range(top_n):
        id_val = I[0][i]
        if id_val >= 0 and D[0][i]<=dis_thred:
            ids.append(id_lst[id_val])
    return ids, I_val, D_val


